import glob
import gzip
import multiprocessing
import os
import shutil
import sys
import tarfile

import numpy as np
from keras.utils import to_categorical

from weiqi.data.generator import DataGenerator
from weiqi.dlgo.board import Board, GameState, Move
from weiqi.dlgo.gotypes import Player, Point
from weiqi.dlgo.sgf.core import SgfGame

from ..dlgo.utils import R as DR
from ..utils import R
from .index_processor import KGSIndex
from .sampling import Sampler


def Worker(jobinfo):
    try:
        clazz, encoder, zip_file, data_filename, game_list = jobinfo
        clazz(encoder_name=encoder).ProcessZip(zip_file, data_filename, game_list)
    except (KeyboardInterrupt, SystemExit):
        raise Exception('>>> Exiting child process.')
    
class GoDataProcessor:
    def __init__(self, encoder_name='simple', data_dir='kgs_data') -> None:
        self.encoder_name = encoder_name
        self.encoder = DR.GetEncoderByName(encoder_name, R.BoardSize)
        self.data_dir = data_dir

    def LoadData(self, data_type='train', num_samples=1000, use_generator=False):
        index = KGSIndex(data_dir=self.data_dir)
        index.DownloadFiles()

        sampler = Sampler(data_dir=self.data_dir)
        data = sampler.DrawData(data_type, num_samples)

        self.MapToWorker(data_type, data)
        if use_generator:
            generator = DataGenerator(self.data_dir, data)
            return generator
        else:
            features_and_labels = self.ConsolidateGames(data_type, data)
            return features_and_labels

    def MapToWorker(self, data_type, samples):
        zip_names = set()
        indices_by_zip_name = {}
        for filename, index in samples:
            zip_names.add(filename)
            if filename not in indices_by_zip_name:
                indices_by_zip_name[filename] = []
            indices_by_zip_name[filename].append(index)

        zips_to_process = []
        for zipname in zip_names:
            base_name = zipname.replace('.tar.gz', '')
            data_filename = base_name + data_type
            if not os.path.isfile(self.data_dir + '/' + data_filename):
                zips_to_process.append((self.__class__, self.encoder_name, zipname, data_filename,
                                        indices_by_zip_name[zipname]))
                
        cores = multiprocessing.cpu_count()
        pool = multiprocessing.Pool(processes=cores)
        p = pool.map_async(Worker, zips_to_process)
        try:
            _ = p.get()
        except KeyboardInterrupt:
            pool.terminate()
            pool.join()
            sys.exit(-1)

    def UnzipData(self, zip_filename):
        this_gz = gzip.open(self.data_dir + '/' + zip_filename)
        tar_filename = zip_filename[0:-3]
        this_tar = open(self.data_dir + '/' + tar_filename, 'wb')
        shutil.copyfileobj(this_gz, this_tar)
        this_tar.close()
        return tar_filename
    
    def ProcessZip(self, zip_filename, data_filename, game_list):
        tar_file = self.UnzipData(zip_filename)
        zip_file = tarfile.open(self.data_dir + '/' + tar_file)
        name_list = zip_file.getnames()
        total_samples = self.NumberTotalExamples(zip_file, game_list, name_list)

        shape = self.encoder.Shape()
        feature_shape = np.insert(shape, 0, np.asarray([total_samples]))
        features = np.zeros(feature_shape)
        labels = np.zeros((total_samples,))

        counter = 0
        for index in game_list:
            name = name_list[index+1]
            if not name.endswith('.sgf'):
                raise ValueError(f'{name} is not a valid sgf')
            sgf_content = zip_file.extractfile(name).read()
            sgf_content = sgf_content.decode('utf-8')
            sgf = SgfGame.FromString(sgf_content)
            game_state, first_move_done = self.GetHandicap(sgf)
            for item in sgf.MainSequenceIter():
                color, move_tuple = item.GetMove()
                point = None 
                if color is not None:
                    if move_tuple is not None:
                        row,col = move_tuple
                        point = Point(row+1, col+1)
                        move = Move.Play(point)
                    else:
                        move = Move.Pass()
                    if first_move_done and point is not None:
                        features[counter] = self.encoder.Encode(game_state)
                        labels[counter] = self.encoder.EncodePoint(point)
                        counter += 1
                    game_state = game_state.ApplyMove(move)
                    first_move_done = True 
        feature_file_base = self.data_dir + '/' + data_filename + '_features_'
        label_file_base = self.data_dir + '/' + data_filename + '_labels_'
        chunk = 0
        chunk_size = 1024
        while features.shape[0] >= chunk_size:
            feature_file = feature_file_base + f'{chunk}'
            label_file = label_file_base + f'{chunk}'
            chunk += 1
            current_features, features = features[:chunk_size], features[chunk_size:]
            current_labels, labels = labels[:chunk_size], labels[chunk_size:]
            np.save(feature_file, current_features)
            np.save(label_file, current_labels)

    def ConsolidateGames(self, data_type, samples):
        files_needed = set(filename for filename, index in samples)
        file_names = []
        for zip_file_name in files_needed:
            file_name = zip_file_name.replace('.tar.gz', '') + data_type
            file_names.append(file_name)
        feature_list = []
        label_list = []
        for filename in file_names:
            file_prefix = filename.replace('.tar.gz', '')
            base = self.data_dir + '/' + file_prefix + '_features_*.npy'
            for feature_file in glob.glob(base):
                label_file = feature_file.replace('features', 'labels')
                x = np.load(feature_file)
                y = np.load(label_file)
                x = x.astype('float32')
                y = to_categorical(y.astype(int), R.BoardSize * R.BoardSize)
                feature_list.append(x)
                label_list.append(y)
        features = np.concatenate(feature_list, axis=0)
        labels = np.concatenate(label_list, axis=0)
        np.save(f'{self.data_dir}/features_{data_type}.npy', features)
        np.save(f'{self.data_dir}/labels_{data_type}.npy', labels)
        return features, labels

    
    @staticmethod
    def GetHandicap(sgf: SgfGame):
        go_board = Board(R.BoardSize, R.BoardSize)
        first_move_done = False
        move = None 
        game_state = GameState.NewGame(R.BoardSize)
        if sgf.GetHandicap() is not None and sgf.GetHandicap() != 0:
            for setup in sgf.GetRoot().GetSetupStones():
                for move in setup:
                    row,col = move 
                    go_board.PlaceStone(Player.black, Point(row+1, col+1))
            first_move_done = True
            game_state = GameState(go_board, Player.white, None, move)
        return game_state, first_move_done

    def NumberTotalExamples(self, zip_file, game_list, name_list):
        total_examples = 0 
        for index in game_list:
            name:str = name_list[index + 1]
            if name.endswith('.sgf'):
                sgf_content = zip_file.extractfile(name).read()
                sgf_content = sgf_content.decode('utf-8')
                sgf = SgfGame.FromString(sgf_content)
                game_state, first_move_done = self.GetHandicap(sgf)
                num_moves = 0
                for item in sgf.MainSequenceIter():
                    color, move = item.GetMove()
                    if color is not None:
                        if first_move_done:
                            num_moves += 1
                        first_move_done = True 
                total_examples = total_examples + num_moves
            else:
                raise ValueError(f'{name} is not a valid sgf')
        return total_examples



